{
    "cells": [
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "# SAHI with RT-DETR for Sliced Inference"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/obss/sahi/blob/main/demo/inference_for_yolov5.ipynb)"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "## 0. Preparation"
            ]
        },
        {
            "attachments": {},
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "- Install latest version of SAHI and ultralytics:"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "!pip install -U torch sahi ultralytics"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "import os\n",
                "\n",
                "os.getcwd()"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "- Import required modules:"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "# arrange an instance segmentation model for test\n",
                "from IPython.display import Image\n",
                "\n",
                "from sahi import AutoDetectionModel\n",
                "from sahi.predict import get_prediction, get_sliced_prediction, predict\n",
                "from sahi.utils.cv import read_image\n",
                "from sahi.utils.file import download_from_url\n",
                "from sahi.utils.rtdetr import download_rtdetrl_model"
            ]
        },
        {
            "attachments": {},
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "- Download a rtdetr model and two test images:"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "# download rtdetr-l model to 'models/rtdetr-l.pt'\n",
                "rtdetr_model_path = \"models/rtdetr-l.pt\"\n",
                "download_rtdetrl_model(rtdetr_model_path)\n",
                "\n",
                "# download test images into demo_data folder\n",
                "download_from_url(\n",
                "    \"https://raw.githubusercontent.com/obss/sahi/main/demo/demo_data/small-vehicles1.jpeg\",\n",
                "    \"demo_data/small-vehicles1.jpeg\",\n",
                ")\n",
                "download_from_url(\n",
                "    \"https://raw.githubusercontent.com/obss/sahi/main/demo/demo_data/terrain2.png\", \"demo_data/terrain2.png\"\n",
                ")"
            ]
        },
        {
            "attachments": {},
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "## 1. Standard Inference with a RTDETR Model"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "- Instantiate a detection model by defining model weight path and other parameters:"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "detection_model = AutoDetectionModel.from_pretrained(\n",
                "    model_type=\"rtdetr\",\n",
                "    model_path=rtdetr_model_path,\n",
                "    confidence_threshold=0.3,\n",
                "    device=\"cpu\",  # or 'cuda:0'\n",
                ")"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "- Perform prediction by feeding the `get_prediction` function with an image path and a DetectionModel instance:"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "result = get_prediction(\"demo_data/small-vehicles1.jpeg\", detection_model)"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "- Or perform prediction by feeding the `get_prediction` function with a numpy image and a DetectionModel instance:"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "result = get_prediction(read_image(\"demo_data/small-vehicles1.jpeg\"), detection_model)"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "- Visualize predicted bounding boxes and masks over the original image:"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "result.export_visuals(export_dir=\"demo_data/\")\n",
                "\n",
                "Image(\"demo_data/prediction_visual.png\")"
            ]
        },
        {
            "attachments": {},
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "## 2. Sliced Inference with a RTDETR Model"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "- To perform sliced prediction we need to specify slice parameters. In this example we will perform prediction over slices of 256x256 with an overlap ratio of 0.2:"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "result = get_sliced_prediction(\n",
                "    \"demo_data/small-vehicles1.jpeg\",\n",
                "    detection_model,\n",
                "    slice_height=256,\n",
                "    slice_width=256,\n",
                "    overlap_height_ratio=0.2,\n",
                "    overlap_width_ratio=0.2,\n",
                ")"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "- Visualize predicted bounding boxes and masks over the original image:"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "result.export_visuals(export_dir=\"demo_data/\")\n",
                "\n",
                "Image(\"demo_data/prediction_visual.png\")"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "## 3. Prediction Result"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "- Predictions are returned as [sahi.prediction.PredictionResult](sahi/prediction.py), you can access the object prediction list as:"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "object_prediction_list = result.object_prediction_list"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "object_prediction_list[0]"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "- ObjectPrediction's can be converted to [COCO annotation](https://cocodataset.org/#format-data) format:"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "result.to_coco_annotations()[:3]"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "- ObjectPrediction's can be converted to [COCO prediction](https://github.com/i008/COCO-dataset-explorer) format:"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "result.to_coco_predictions(image_id=1)[:3]"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "- ObjectPrediction's can be converted to [imantics](https://github.com/jsbroks/imantics) annotation format:"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "result.to_imantics_annotations()[:3]"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "- ObjectPrediction's can be converted to [fiftyone](https://github.com/voxel51/fiftyone) detection format:"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "result.to_fiftyone_detections()[:3]"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "## 4. Batch Prediction"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "- Set model and directory parameters:"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "model_type = \"rtdetr\"\n",
                "model_path = rtdetr_model_path\n",
                "model_device = \"cpu\"  # or 'cuda:0'\n",
                "model_confidence_threshold = 0.4\n",
                "\n",
                "slice_height = 256\n",
                "slice_width = 256\n",
                "overlap_height_ratio = 0.2\n",
                "overlap_width_ratio = 0.2\n",
                "\n",
                "source_image_dir = \"demo_data/\""
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "- Perform sliced inference on given folder:"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "predict(\n",
                "    model_type=model_type,\n",
                "    model_path=model_path,\n",
                "    model_device=model_device,\n",
                "    model_confidence_threshold=model_confidence_threshold,\n",
                "    source=source_image_dir,\n",
                "    slice_height=slice_height,\n",
                "    slice_width=slice_width,\n",
                "    overlap_height_ratio=overlap_height_ratio,\n",
                "    overlap_width_ratio=overlap_width_ratio,\n",
                ")"
            ]
        }
    ],
    "metadata": {
        "kernelspec": {
            "display_name": "test",
            "language": "python",
            "name": "python3"
        },
        "language_info": {
            "codemirror_mode": {
                "name": "ipython",
                "version": 3
            },
            "file_extension": ".py",
            "mimetype": "text/x-python",
            "name": "python",
            "nbconvert_exporter": "python",
            "pygments_lexer": "ipython3",
            "version": "3.10.9"
        },
        "vscode": {
            "interpreter": {
                "hash": "244b47d5824a96a4079632e50977464d968e13d2c337f65c905f8da81a0b4f95"
            }
        }
    },
    "nbformat": 4,
    "nbformat_minor": 4
}
